#include "DINO.h"
#include "llvm/Analysis/Dominators.h"
#include "llvm/Support/InstIterator.h"
#include "llvm/IR/Module.h"
#include "llvm/Support/raw_ostream.h"
#include <cstdio>

using namespace llvm;

DINOStoreAnalysis::DINOStoreAnalysis ()
    : ModulePass(ID)
{ }

DINOStoreAnalysis::~DINOStoreAnalysis ()
{ }

bool DINOStoreAnalysis::doInitialization (Module &M)
{
    return false;
}

// Recognize stores to nonvolatile locations.
bool DINOStoreAnalysis::runOnModule (Module &M)
{
    // Bookkeeping: note which global variables are marked with a prefix
    // indicating they are nonvolatile, and simply record this fact in a
    // module-wide map.
    for (auto &G : M.getGlobalList()) {
        StringRef varName = G.getName();
        if (varName.find(NVGlobalPrefix) == 0) {
            GlobalNVness[varName] = true;
        }
    }

    // For each Instruction in each Function, note whether it's a store to
    // a variable we just marked as nonvolatile.
    for (auto &F : M.getFunctionList()) {
        NVVars[&F] = DINOStoreAnalysis::GlobalVarSet();
        getNVVars(F, NVVars[&F]);
    }

    return false;
}

/* Returns an InstList of StoreInstructions in F that touch NV variables. */
DINOStoreAnalysis::InstList *DINOStoreAnalysis::getNVStores (Function &F)
{
    if (NVStores.count(&F))
        return &NVStores[&F];

    NVStores[&F] = InstList();
    InstList *IL = &NVStores[&F];

    for (inst_iterator I = inst_begin(F), E = inst_end(F); I != E; ++I) {
        if (isNVStore(*I)) {
            IL->push_back(&*I);
        }
    }

    outs() << "Function " << F.getName() << " has " << IL->size() <<
        " NV store(s)\n";
    return IL;
}

GlobalVariable *DINOStoreAnalysis::getVarFromStore (StoreInst &SI)
{
    return dyn_cast<GlobalVariable>(SI.getOperand(1));
}

/* Populates GVS with the list of NV variables this function may touch. */
void DINOStoreAnalysis::getNVVars (Function &F, DINOStoreAnalysis::GlobalVarSet &GVS)
{
    DINOStoreAnalysis::InstList &IL = *getNVStores(F);
    for (auto &I : IL) {
        if (StoreInst *SI = dyn_cast<StoreInst>(I)) {
            if (GlobalVariable *GV = getVarFromStore(*SI)) {
                GVS.push_back(GV);
            }
        }
    }

    if (!GVS.size())
        return;

    outs() << "NV stores in function " << F.getName() << ":\n";
    for (auto &S : GVS) {
        outs() << "  * " << *S << "\n";
    }
}

/* Returns true if I stores to a nonvolatile variable. */
bool DINOStoreAnalysis::isNVStore (Instruction &I)
{
    StoreInst *SI = dyn_cast<StoreInst>(&I);
    if (!SI)
        return false;

    const GlobalVariable *GV = getVarFromStore(*SI);
    if (!GV)
        return false;

    return (GlobalNVness.count(GV->getName()) > 0);
}

// Returns true if I is a DINO task boundary.
bool DINOStoreAnalysis::isTaskBoundary (Instruction &I)
{
    CallInst *call = dyn_cast<CallInst>(&I);
    if (!call)
        return false;

    Function *func = call->getCalledFunction();
    if (!func)
        return false;

    return func->getName() == DINOTaskBoundaryFunctionName;
}

/**
 * Starting with an Instruction, traverse the CFG backward on all possible
 * paths, stopping on each path when we hit a task boundary.  Return the set of
 * these task boundaries.
 */
DINOStoreAnalysis::InstList DINOStoreAnalysis::getTaskBoundaries (Instruction &I)
{
    DINOStoreAnalysis::InstList TBs;

    // TODO write me

    return TBs;
}

void DINOStoreAnalysis::getAnalysisUsage (AnalysisUsage &AU) const
{
    AU.setPreservesAll(); // this is an analysis pass
}

bool DINOStoreAnalysis::doFinalization (Module &M)
{
    return false;
}

const char *DINOStoreAnalysis::getPassName () const {
    return "DINO Store Analysis";
}

char DINOStoreAnalysis::ID = 0;
